clear; close all;

% Parameters: recency/similarity/salience
mu = 0.0 * 100;
sigma = 0.02 * 100;
sigma_r = 0.02;
sigma_s = 5;
sigma_a = 0.05;

% Past History
t = 119;
T = t + 1;

% Set a fixed seed for reproducibility
rng(42);

% Generate fixed random historical return
rt = normrnd(mu, sigma, [t, 1]);

% Generate rT vector with 21 equally spaced values between -10 and 10 rT = linspace(-10, 10, 21);
rTs = linspace(-5, 5, 21);
vec = (t:-1:0)';

for i = 1:length(rTs)
    rT = rTs(i);
    rt(t + 1) = rT;

    % similarity
    s_rt = exp(- (rt - rT) .^ 2 / (2 * sigma_s ^ 2));
    s_rt_norm(i, :) = s_rt ./ sum(s_rt);

    % recency
    r_rt = exp(-sigma_r * vec);
    r_rt_norm(i, :) = r_rt / sum(r_rt);

    % salience
    a_rt = exp((rt - mu) .^ 2 * sigma_a);
    a_rt_norm(i, :) = a_rt / sum(a_rt);

    %scatter(rt, sa_rt_norm);

    % retrival probability and normalization
    pt = s_rt_norm(i, :) .* r_rt_norm(i, :) .* a_rt_norm(i, :);
    pt_norm(i, :) = pt ./ sum(pt);
    recall_ret(i) = sum(rt .* pt_norm(i, :)');

end

% Plots
rT = 2;
t = t + 1;

figure;
plot(1:t, rt, '-', 1:t, ones(t, 1) * 2, '--', 'LineWidth', 1.25)
xlabel('Time', 'Interpreter', 'latex')
ylabel('Return (\%)', 'Interpreter', 'latex')
legend('Return $r_t$', 'Cue $r_T$', 'Location', 'northwest', 'Interpreter', 'latex')
grid
exportgraphics(gcf, '../Figures/return.pdf', 'Resolution', 300)

%
K = 15;
rTs(K);

figure;
plot(1:t, pt_norm(K, :) * 100, '-', 1:t, s_rt_norm(K, :) * 100, '--', 1:t, r_rt_norm(K, :) * 100, ':', 1:t, a_rt_norm(K, :) * 100, '-.', 'LineWidth', 1.25)
xlabel('Time', 'Interpreter', 'latex')
ylabel('Probability (\%)', 'Interpreter', 'latex')
legend('Probability', 'Similarity', 'Recency', 'Salience', 'Location', 'northwest', 'Interpreter', 'latex')
grid
xlim([1, t])
exportgraphics(gcf, '../Figures/recallprob.pdf', 'Resolution', 300)

figure;
plot(rTs, recall_ret, 'LineWidth', 1.25)
xlabel('Cue (\%)', 'Interpreter', 'latex')
ylabel('Expected Return (\%)', 'Interpreter', 'latex')
grid
exportgraphics(gcf, '../Figures/cuedrecall.pdf', 'Resolution', 300)
